#! /usr/bin/env python3
import rospy
from nav_msgs.msg import Odometry
import json
import sys

drone_info_dict = {}
start_time = 0
last_time = 0
drone_num = int(sys.argv[1])

def odom_callback(msg, drone_id):
    global start_time, last_time
    cur_time = rospy.get_time()
    if round((cur_time - last_time), 1) < 0.1:
      return
    x = msg.pose.pose.position.x
    y = msg.pose.pose.position.y
    z = msg.pose.pose.position.z
    yaw = msg.pose.pose.orientation.z
    t = cur_time - start_time
    last_time = cur_time
    
    if drone_id not in drone_info_dict:
        drone_info_dict[drone_id] = []
    drone_info_dict[drone_id].append({'x': round(x, 2), 'y': round(y, 2), 'z': round(z, 2), 'yaw': round(yaw, 2), 't': round(t, 2)})
    print(f'current time: {t}')

def save_drone_info():
    # 将字典保存为JSON文件
    with open(f'drone{drone_num}_info.json', 'w') as file:
        json.dump(drone_info_dict, file)

if __name__ == '__main__':
    rospy.init_node('drone_info_subscriber')
    start_time = rospy.get_time()
    last_time = rospy.get_time()
    drone_ids = ['drone_{}'.format(i) for i in range(drone_num)]
    drone_topics = ['/drone_{}_visual_slam/odom'.format(i) for i in range(drone_num)]

    subscribers = []
    for drone_topic in drone_topics:
        sub = rospy.Subscriber(drone_topic, Odometry, odom_callback, callback_args=drone_topic)
        subscribers.append(sub)

    rate = rospy.Rate(10)

    while not rospy.is_shutdown():
        # print(drone_info_dict)
        save_drone_info()
        rate.sleep()